import cv2
import numpy as np
from paddleocr import PaddleOCR
import os
from datetime import datetime


def video_text_recognition(video_path, output_dir="ocr_results"):
    """
    使用PaddleOCR识别视频中的文字
    参数:
        video_path: 视频文件路径
        output_dir: 结果保存目录
    """
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 初始化PaddleOCR，使用中英文模型
    # 修正参数：移除show_log，使用新的方向检测参数
    ocr = PaddleOCR(use_textline_orientation=True, lang="ch")

    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频文件: {video_path}")
        return

    # 获取视频基本信息
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = frame_count / fps if fps > 0 else 0

    print(f"视频信息:")
    print(f"  帧率: {fps:.2f} FPS")
    print(f"  总帧数: {frame_count}")
    print(f"  时长: {duration:.2f} 秒")

    # 记录结果的文件
    result_file = os.path.join(output_dir, f"ocr_results.txt")

    frame_idx = 0
    success = True

    with open(result_file, "w", encoding="utf-8") as f:
        f.write(f"视频文字识别结果 - {video_path}\n")
        f.write(f"识别时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"帧率: {fps:.2f} FPS, 总帧数: {frame_count}, 时长: {duration:.2f} 秒\n\n")
        process_interval = 30
        frame_last = np.array([])
        frame_last_judge = None
        while success:
            # 读取一帧
            success, frame = cap.read()
            if not success:
                break
            # 计算当前时间点
            current_time = frame_idx / fps if fps > 0 else 0
            color = ""
            if frame_idx % process_interval == 0:
                print(f"处理帧 {frame_idx}/{frame_count} ({current_time:.2f}秒)")
                img_save_dir = "saved_frames"  #图像保存目录
                os.makedirs(img_save_dir, exist_ok=True)  # 自动创建目录，已存在则不报错

                # 2. 生成带格式的文件名（包含帧索引，确保唯一性）
                img_filename = f"frame_{frame_idx}.jpg"  # 例如：frame_0.jpg、frame_5.jpg
                img_path = os.path.join(img_save_dir, img_filename)
                cv2.imwrite(img_path, frame)
                img1_cut = frame[930:1090,650:790]
                img2_cut = frame[720:850,1880:2240]
                # 识别文字
                result1 = ocr.predict(img1_cut)
                result2 = ocr.predict(img2_cut)
                frame_now = frame[210:720,1880:2240].astype(np.int16)#注意要转成int16，否则会溢出
                if frame_last_judge != None:
                    #print("形状：",frame_last.shape)
                    for i in range(frame_last.shape[0]):
                        if color != "":
                            break
                        for j in range(frame_last.shape[1]):
                            if color != "":
                                break
                            if abs(frame_last[i][j][0]-frame_now[i][j][0])<20 and abs(frame_last[i][j][1]-frame_now[i][j][1])<20 and abs(frame_last[i][j][2]-frame_now[i][j][2])<20:
                                continue
                            else:
                                if (frame_now[i][j][1]-frame_last[i][j][1]>180) and (frame_now[i][j][2]-frame_last[i][j][2]>180) and (frame_now[i][j][0]-frame_last[i][j][0]<120):
                                    color = 'y'
                                    print(color)
                                elif (frame_now[i][j][1]-frame_last[i][j][1]>180) and (frame_now[i][j][2]-frame_last[i][j][2]<120):
                                    color = 'g'
                                    print(color)
                                elif (frame_now[i][j][2]-frame_last[i][j][2]>180) and (frame_now[i][j][1]-frame_last[i][j][1]<120):
                                    color = 'r'
                                    print(color)
                frame_last = frame_now
                frame_last_judge = True
                # 处理识别结果
                if result1 is not None and len(result1) > 0:
                    for i in range(len(result1)):
                        line1 = result1[i]
                        line2 = result2[i]
                        if line1 is not None:
                            for word_info in line1:
                                texts = line1['rec_texts']+line2['rec_texts']

                    if texts:
                        #保存
                        time_str = f"[{current_time:.2f}秒] "
                        text_str = " ".join(texts)
                        f.write(time_str + text_str + ' ' +color + "\n")
                        if color != "":
                            print(time_str + text_str + f"新颜色： {color}")
                        else:
                            print(time_str + text_str)

            frame_idx += 1
    # 释放资源
    cap.release()
    print(f"\n识别完成！结果已保存到: {result_file}")
    return result_file


if __name__ == "__main__":
    video_path = "./input_video.mp4"
    video_text_recognition(video_path)